package org.apache.flink.connector.scylladb.source;

import com.datastax.driver.core.*;
import com.datastax.driver.core.policies.DCAwareRoundRobinPolicy;
import com.datastax.driver.core.policies.DefaultRetryPolicy;
import com.datastax.driver.core.policies.TokenAwarePolicy;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.scylladb.serialization.ScyllaDBRowDataDeserializer;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.types.logical.RowType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import java.time.Duration;
import java.util.Iterator;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * ScyllaDB 源函数，负责从 ScyllaDB 读取数据
 */
public class ScyllaDBSourceFunction extends RichSourceFunction<RowData> {

    private static final Logger LOG = LoggerFactory.getLogger(ScyllaDBSourceFunction.class);

    private final String url;
    private final String keySpace;
    private final String tableName;
    private final String username;
    private final String password;
    private final int fetchSize;
    private final String partitionColumn;
    private final int connectionPoolLocalSize;
    private final int connectionPoolRemoteSize;
    private final Duration connectionTimeout;
    private final Duration requestTimeout;
    private final String consistencyLevel;
    private final boolean sslEnabled;
    private final String sslKeystorePath;
    private final String sslKeystorePassword;
    private final String sslTruststorePath;
    private final String sslTruststorePassword;
    private final RowType rowType;
    private final int[][] projectedFields;
    private final long limit;

    private transient Cluster cluster;
    private transient Session session;
    private transient ScyllaDBRowDataDeserializer deserializer;
    private final AtomicBoolean isRunning = new AtomicBoolean(true);

    public ScyllaDBSourceFunction(
            String url,
            String keySpace,
            String tableName,
            String username,
            String password,
            int fetchSize,
            String partitionColumn,
            int connectionPoolLocalSize,
            int connectionPoolRemoteSize,
            Duration connectionTimeout,
            Duration requestTimeout,
            String consistencyLevel,
            boolean sslEnabled,
            String sslKeystorePath,
            String sslKeystorePassword,
            String sslTruststorePath,
            String sslTruststorePassword,
            RowType rowType,
            int[][] projectedFields,
            long limit) {
        
        this.url = url;
        this.keySpace = keySpace;
        this.tableName = tableName;
        this.username = username;
        this.password = password;
        this.fetchSize = fetchSize;
        this.partitionColumn = partitionColumn;
        this.connectionPoolLocalSize = connectionPoolLocalSize;
        this.connectionPoolRemoteSize = connectionPoolRemoteSize;
        this.connectionTimeout = connectionTimeout;
        this.requestTimeout = requestTimeout;
        this.consistencyLevel = consistencyLevel;
        this.sslEnabled = sslEnabled;
        this.sslKeystorePath = sslKeystorePath;
        this.sslKeystorePassword = sslKeystorePassword;
        this.sslTruststorePath = sslTruststorePath;
        this.sslTruststorePassword = sslTruststorePassword;
        this.rowType = rowType;
        this.projectedFields = projectedFields;
        this.limit = limit;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        
        LOG.info("Opening ScyllaDB source connection to {}", url);
        
        // 创建集群连接
        this.cluster = createCluster();
        this.session = cluster.connect(keySpace);
        
        // 创建反序列化器
        this.deserializer = new ScyllaDBRowDataDeserializer(rowType, projectedFields);
        
        LOG.info("ScyllaDB source connection opened successfully");
    }

    @Override
    public void run(SourceContext<RowData> ctx) throws Exception {
        // 构建查询语句
        String query = buildQuery();
        LOG.debug("Executing query: {}", query);

        Statement statement = new SimpleStatement(query);
        statement.setFetchSize(fetchSize);
        statement.setConsistencyLevel(ConsistencyLevel.valueOf(consistencyLevel));
        
        ResultSet resultSet = session.execute(statement);
        Iterator<Row> iterator = resultSet.iterator();
        
        long count = 0;
        while (isRunning.get() && iterator.hasNext()) {
            Row row = iterator.next();
            RowData rowData = deserializer.deserialize(row);
            
            ctx.collect(rowData);
            count++;
            
            // 如果设置了限制，检查是否达到限制
            if (limit > 0 && count >= limit) {
                LOG.debug("Reached limit of {} rows, stopping", limit);
                break;
            }
            
            // 检查是否被取消
            if (!isRunning.get()) {
                LOG.debug("Source function cancelled, stopping");
                break;
            }
        }
        
        LOG.debug("ScyllaDB source finished reading {} rows", count);
    }

    @Override
    public void cancel() {
        LOG.info("Cancelling ScyllaDB source function");
        isRunning.set(false);
    }

    @Override
    public void close() throws Exception {
        LOG.info("Closing ScyllaDB source connection");
        
        if (session != null) {
            session.close();
        }
        
        if (cluster != null) {
            cluster.close();
        }
        
        super.close();
        LOG.info("ScyllaDB source connection closed");
    }

    private Cluster createCluster() {
        String[] hostAndPort = url.split(":");
        String host = hostAndPort[0];
        int port = hostAndPort.length > 1 ? Integer.parseInt(hostAndPort[1]) : 9042;
        
        Cluster.Builder builder = Cluster.builder()
                .addContactPoint(host)
                .withPort(port)
                .withCredentials(username, password)
                .withLoadBalancingPolicy(
                        new TokenAwarePolicy(DCAwareRoundRobinPolicy.builder().build()))
                .withRetryPolicy(DefaultRetryPolicy.INSTANCE)
                .withPoolingOptions(new PoolingOptions()
                        .setConnectionsPerHost(HostDistance.LOCAL, connectionPoolLocalSize, connectionPoolLocalSize)
                        .setConnectionsPerHost(HostDistance.REMOTE, connectionPoolRemoteSize, connectionPoolRemoteSize))
                .withSocketOptions(new SocketOptions()
                        .setConnectTimeoutMillis((int) connectionTimeout.toMillis())
                        .setReadTimeoutMillis((int) requestTimeout.toMillis()));

        // SSL 配置
        if (sslEnabled) {
            try {
                SSLOptions sslOptions = RemoteEndpointAwareJdkSSLOptions.builder()
                        .withSSLContext(createSSLContext())
                        .build();
                builder.withSSL(sslOptions);
            } catch (Exception e) {
                throw new RuntimeException("Failed to configure SSL", e);
            }
        }

        return builder.build();
    }

    private SSLContext createSSLContext() throws Exception {
        // 简化的 SSL 上下文创建，实际使用时需要根据具体需求配置
        return SSLContext.getDefault();
    }

    private String buildQuery() {
        StringBuilder sb = new StringBuilder();
        
        // 构建 SELECT 子句
        sb.append("SELECT ");
        if (projectedFields != null && projectedFields.length > 0) {
            // 如果有投影，只选择需要的列
            for (int i = 0; i < projectedFields.length; i++) {
                if (i > 0) {
                    sb.append(", ");
                }
                String fieldName = rowType.getFieldNames().get(projectedFields[i][0]);
                sb.append(fieldName);
            }
        } else {
            // 选择所有列
            sb.append("*");
        }
        
        sb.append(" FROM ").append(keySpace).append(".").append(tableName);
        
        // 如果指定了分区列，可以添加WHERE子句来优化查询
        if (partitionColumn != null && !partitionColumn.isEmpty()) {
            // 这里可以根据具体需求添加分区过滤逻辑
            // 例如：WHERE partition_column IN (token_range)
        }
        
        // 添加 LIMIT 子句
        if (limit > 0) {
            sb.append(" LIMIT ").append(limit);
        }
        
        return sb.toString();
    }
} 